import wx
import os
import pickle
import numpy as np
from load_data.loadnpz import loadnpz
from Processor import train_model


class TrainModelWindow(wx.Dialog):
    def __init__(self, parent, title):
        super(TrainModelWindow, self).__init__(parent, title=title)
        self.saveModelPath = parent.config_data.trainModelPath
        self.dataNum = parent.trainDataNum.GetValue()
        self.TrainDataPathCtrl = list(range(self.dataNum))
        self.init_ui()

    def init_ui(self):
        dataWildcard = "npz Data File (.npz)" + "|*.npz"
        panel = wx.Panel(self)
        grid_sizer1 = wx.FlexGridSizer(cols=2, vgap=1, hgap=1)
        for i in range(self.dataNum):
            labelText = '选择数据' + str(i+1) + '：'
            label = wx.StaticText(panel, label=labelText)
            grid_sizer1.Add(label, 0, wx.ALL | wx.ALIGN_CENTER_VERTICAL, 5)
            self.TrainDataPathCtrl[i] = wx.FilePickerCtrl(panel, wildcard=dataWildcard, size=(350, 27))
            self.TrainDataPathCtrl[i].SetInitialDirectory(os.path.dirname(self.saveModelPath))
            self.TrainDataPathCtrl[i].GetPickerCtrl().SetLabel('浏览')
            grid_sizer1.Add(self.TrainDataPathCtrl[i], 0, wx.ALL | wx.ALIGN_CENTER_VERTICAL, 5)
        self.SetSize(470, 60*(self.dataNum+1))
        self.Centre()
        grid_sizer2 = wx.FlexGridSizer(cols=2, vgap=1, hgap=1)
        self.TrainModelBtn = wx.Button(panel, label='模型训练开始', size=wx.Size(100, 27))
        self.TrainModelBtn.Bind(wx.EVT_BUTTON, self.on_train_model)
        grid_sizer2.Add(self.TrainModelBtn, 0, wx.ALL | wx.ALIGN_CENTER_VERTICAL, 5)
        self.statusLabel = wx.StaticText(panel, label=' ')
        grid_sizer2.Add(self.statusLabel, 0, wx.ALL | wx.ALIGN_CENTER_VERTICAL, 5)

        grid_sizer = wx.FlexGridSizer(cols=1, vgap=1, hgap=1)
        grid_sizer.Add(grid_sizer1, 0, wx.ALL, 5)
        grid_sizer.Add(grid_sizer2, 0, wx.ALL, 5)
        panel.SetSizerAndFit(grid_sizer)
        panel.Center()
        self.Fit()

    def on_close(self, event):
        self.Close()  # 关闭窗体

    def gather_data(self):
        data_num = len(self.TrainDataPathCtrl)
        train_data_path = list(range(data_num))
        data_x_gather, data_y_gather = None, None
        for i in range(data_num):
            train_data_path[i] = self.TrainDataPathCtrl[i].GetPath()
            if train_data_path[i] == '':
                train_data_path.remove(train_data_path[i])
        train_data_path = list(set(train_data_path))
        for i in range(len(train_data_path)):
            data_x, data_y = loadnpz(train_data_path[i])  # x:(sample, channal, trial)  y:(trial,)
            if i == 0:
                data_x_gather = data_x
                data_y_gather = data_y
            else:
                data_x_gather = np.concatenate((data_x_gather, data_x), axis=2)
                data_y_gather = np.concatenate((data_y_gather, data_y), axis=0)
        return data_x_gather, data_y_gather

    def on_train_model(self, event):
        data_x, data_y = self.gather_data()
        csp_proj_matrix, classifier_model = train_model(data_x, data_y, classifier_type='lda', m=3)
        # TrainModelPath = time.strftime(TrainModelPath + "\\TrainModel_%Y_%m_%d_%H_%M_%S.pkl")
        f1 = open(self.saveModelPath, 'wb')
        pickle.dump(csp_proj_matrix, f1)
        pickle.dump(classifier_model, f1)
        f1.close()
        self.statusLabel.SetLabel('模型训练完成。')



